import pymysql
from config.config import get_mysql_config


class DB:
    """改进后的数据库操作类"""

    def __init__(self):
        mysql_config = get_mysql_config()
        self.host = mysql_config['host']
        self.port = int(mysql_config['port'])
        self.user = mysql_config['user']
        self.password = mysql_config['password']
        self.database = mysql_config['database']
        self.conn = None
        self.cursor = None

    def __enter__(self):
        """支持with语句上下文管理"""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """支持with语句自动关闭连接"""
        self.close()

    def connect(self):
        """建立数据库连接"""
        try:
            self.conn = pymysql.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database,
                charset="utf8mb4",
                cursorclass=pymysql.cursors.DictCursor  # 返回字典格式结果
            )
            self.cursor = self.conn.cursor()
            return True
        except Exception as e:
            raise ConnectionError(f"数据库连接失败: {e}")

    def close(self):
        """安全关闭连接"""
        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()
        self.cursor = None
        self.conn = None

    def execute(self, sql, params=None):
        """执行SQL语句"""
        try:
            self.cursor.execute(sql, params or ())
            self.conn.commit()
            return self.cursor.rowcount
        except Exception as e:
            self.conn.rollback()
            raise RuntimeError(f"SQL执行失败: {e}")

    def query(self, sql, params=None):
        """执行查询并返回结果"""
        try:
            self.cursor.execute(sql, params or ())
            return self.cursor.fetchall()
        except Exception as e:
            raise RuntimeError(f"查询失败: {e}")

    def query_one(self, sql, params=None):
        """执行查询并返回单条结果"""
        try:
            self.cursor.execute(sql, params or ())
            return self.cursor.fetchone()
        except Exception as e:
            raise RuntimeError(f"查询失败: {e}")

    def query_value(self, sql, params=None):
        """查询单个值并确保返回整数"""
        try:
            self.cursor.execute(sql, params or ())
            result = self.cursor.fetchone()
            if result:
                # 处理DictCursor返回的字典结果
                value = list(result.values())[0] if isinstance(result, dict) else result[0]
                return int(value) if value is not None else 0
            return 0
        except Exception as e:
            print(f"查询失败: {str(e)}")
            return 0